const WebSocket = require('ws');
const http = require('http');
const fs = require('fs');
const path = require('path');

// 使用 node-pty 替代 child_process.spawn
const pty = require('node-pty');

const PORT = 8040;
const MAX_CONNECTIONS = 10;
const DEFAULT_COLS = 80;
const DEFAULT_ROWS = 24;
const MIN_COLS = 40;
const MIN_ROWS = 10;
const MAX_COLS = 500;
const MAX_ROWS = 200;

let connectionCount = 0;

// HTTP 服务器（用于提供 HTML）
const server = http.createServer((req, res) => {
    if (req.url === '/' || req.url === '/index.html') {
        const htmlPath = path.join(__dirname, 'index.html');
        fs.readFile(htmlPath, 'utf8', (err, data) => {
            if (err) {
                console.error('HTML file not found:', err.message);
                res.writeHead(500);
                res.end('HTML file not found. Please create index.html');
            } else {
                res.writeHead(200, { 'Content-Type': 'text/html' });
                res.end(data);
            }
        });
    } else {
        res.writeHead(404);
        res.end();
    }
});

const wss = new WebSocket.Server({ server });

wss.on('connection', (ws, req) => {
    const clientIp = req.headers['x-forwarded-for'] || req.connection.remoteAddress;
    console.log(`New connection from ${clientIp}`);
    
    if (connectionCount >= MAX_CONNECTIONS) {
        console.log(`Connection limit reached (${MAX_CONNECTIONS}), rejecting ${clientIp}`);
        ws.close(1013, 'Server busy: Maximum connections reached');
        return;
    }
    
    connectionCount++;
    console.log(`Active connections: ${connectionCount}/${MAX_CONNECTIONS}`);

    // ✅ 使用 node-pty 创建伪终端
    const shell = process.platform === 'win32' ? 'powershell.exe' : 'zsh';
    const ptyProcess = pty.spawn(shell, [], {
        name: 'xterm-256color',
        cols: DEFAULT_COLS,
        rows: DEFAULT_ROWS,
        cwd: process.env.HOME || process.env.USERPROFILE,
        env: process.env
    });

    // 终端输出 → 浏览器
    ptyProcess.onData((data) => {
        if (ws.readyState === WebSocket.OPEN) {
            ws.send(data);
        }
    });

    // 浏览器输入 → 终端
    ws.on('message', (data) => {
        try {
            // 检查是否为 resize 消息
            const message = JSON.parse(data);
            if (message.type === 'resize' && typeof message.cols === 'number' && typeof message.rows === 'number') {
                // 验证尺寸范围
                const cols = Math.max(MIN_COLS, Math.min(MAX_COLS, Math.floor(message.cols)));
                const rows = Math.max(MIN_ROWS, Math.min(MAX_ROWS, Math.floor(message.rows)));
                
                // 只有当尺寸真正改变时才执行 resize
                if (cols !== ptyProcess.cols || rows !== ptyProcess.rows) {
                    ptyProcess.resize(cols, rows);
                    console.log(`PTY resized to ${cols}x${rows} for ${clientIp}`);
                }
                return;
            }
        } catch (e) {
            // 不是 JSON 格式或无效的 resize 消息，当作普通输入处理
            if (e instanceof SyntaxError) {
                // 普通文本输入
            } else {
                console.warn('Invalid resize message:', e.message);
            }
        }
        
        // 普通终端输入
        if (ptyProcess && ptyProcess.write) {
            ptyProcess.write(data);
        }
    });

    // 终端退出
    ptyProcess.onExit((e) => {
        console.log(`PTY exited: code=${e.exitCode}, signal=${e.signal} for ${clientIp}`);
        if (ws.readyState === WebSocket.OPEN) {
            ws.close(1000, `Shell exited: ${e.exitCode}`);
        }
    });

    // WebSocket 关闭
    ws.on('close', (code, reason) => {
        console.log(`WebSocket closed: ${code} - ${reason} for ${clientIp}`);
        connectionCount--;
        console.log(`Active connections: ${connectionCount}/${MAX_CONNECTIONS}`);
        if (ptyProcess && !ptyProcess.killed) {
            try {
                ptyProcess.kill();
                console.log(`PTY process killed for ${clientIp}`);
            } catch (err) {
                console.error(`Failed to kill PTY process: ${err.message}`);
            }
        }
    });

    ws.on('error', (err) => {
        console.error(`WebSocket error for ${clientIp}:`, err.message);
        if (ptyProcess && !ptyProcess.killed) {
            try {
                ptyProcess.kill();
            } catch (err) {
                console.error(`Failed to kill PTY process on error: ${err.message}`);
            }
        }
    });
});

// 错误处理
wss.on('error', (error) => {
    console.error('WebSocket server error:', error);
});

server.on('error', (error) => {
    console.error('HTTP server error:', error);
});

// 优雅关闭
process.on('SIGINT', () => {
    console.log('\nShutting down server gracefully...');
    wss.clients.forEach(client => {
        if (client.readyState === WebSocket.OPEN) {
            client.close(1001, 'Server shutting down');
        }
    });
    
    // 等待所有连接关闭
    setTimeout(() => {
        server.close(() => {
            console.log('Server closed');
            process.exit(0);
        });
    }, 2000);
});

server.listen(PORT, () => {
    console.log(`🚀 Terminal server running on http://localhost:${PORT}`);
    console.log(`🔗 WebSocket server running on ws://localhost:${PORT}`);
    console.log(`👥 Maximum connections: ${MAX_CONNECTIONS}`);
});
